# -*- coding: utf-8 -*-
"""
1. img -> JPEGImages
   xml -> Annotations

2. 将子目录中的图像和标注文件汇集到一个标准VOC标准目录中

原来的目录组织形式如下：

plane
    |
    |---J20
           |--J20_1 (xml, img)
           |--J20_2 (xml, img)
    |
    |
    |---J10
           |--J10_1 (xml, img)
           |--J10_2 (xml，img)

整理后：

plane
    |
    |--Annotations
    |--JPEGImages

compare to merge_subVOC.py:

change object class labels according to directory name.
目标类别同一命名成跟目录相同的大类。例如：原来类别为J10, J20，存在‘plane’目录下， 在脚本运行完成以后，xml文件中的‘class’统一变为‘plane’
"""

import os
import shutil
import sys
import argparse
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def reorganise_directory(root_dir, save_dir, rename=False):
    img_count = 0
    xml_count = 0
    object_classes = os.listdir(root_dir)
    for object_class in object_classes:
        for root, dirs, files in os.walk(os.path.join(root_dir, object_class)):
            if files:
                old_dir_name = root
                new_dir_name = None
                if rename:
                    if os.path.split(old_dir_name)[-1] == "xml":
                        new_dir_name = old_dir_name.replace(os.path.split(old_dir_name)[-1], "Annotations")
                    elif os.path.split(old_dir_name)[-1] == "img":
                        new_dir_name = old_dir_name.replace(os.path.split(old_dir_name)[-1], "JPEGImages")
                    if new_dir_name is not None:
                        os.rename(old_dir_name, new_dir_name)
                for file in files:
                    old_file_full_path = os.path.join(root, file)
                    file_path = root
                    if len(file_path.split("/")) >= 3:
                        if file_path.split("/")[-3] is not None and file_path.split("/")[-2] is not None:
                            save_file_name = file_path.split("/")[-3] + "_" + file_path.split("/")[-2] + "_" + file
                        elif file_path.split("/")[-2] is not None:
                            save_file_name = file_path.split("/")[-2] + "_" + file
                        else:
                            save_file_name = file
                    # if len(save_file_name.split(" ")) >= 2:
                    #     save_file_name = ''.join(save_file_name.split())
                    if os.path.split(old_dir_name)[-1] == "img" or os.path.split(old_dir_name)[-1] == "JPEGImages":
                        # new_file_full_path = os.path.join(save_dir, "JPEGImages", ''.join(save_file_name.split()))
                        new_file_full_path = os.path.join(save_dir, "JPEGImages", save_file_name)
                        img_count += 1
                    elif os.path.split(old_dir_name)[-1] == "xml" or os.path.split(old_dir_name)[-1] == "Annotations":
                        xml_file_full_path = old_file_full_path
                        tree = ET.parse(xml_file_full_path)
                        xml_root = tree.getroot()
                        for obj in xml_root.iter('object'):
                            object_class_in_xml = obj.find('name').text
                            obj.find('name').text = object_class
                        if not os.path.exists(os.path.join(save_dir, "Annotations")):
                            os.makedirs(os.path.join(save_dir, "Annotations"))
                        new_file_full_path = os.path.join(save_dir, "Annotations", save_file_name)
                        tree.write(new_file_full_path)
                        xml_count += 1
                    if new_file_full_path is not None:
                        if not os.path.exists(os.path.split(new_file_full_path)[0]):
                            os.makedirs(os.path.split(new_file_full_path)[0])

                    if os.path.exists(new_file_full_path):
                        continue
                    shutil.copyfile(old_file_full_path, new_file_full_path)
                    if img_count % 1000 == 0:
                        print("copied images: {}".format(img_count))
                    if xml_count % 1000 == 0:
                        print("copied xml: {}".format(xml_count))

    print("*"*80)
    print("Total img copied: {}".format(img_count))
    print("Total xml copied: {}".format(xml_count))


if __name__ == '__main__':

    PARSER = argparse.ArgumentParser(description="arguments")
    PARSER.add_argument("--root_dir", type=str,
                        default="/zongshiban_data/2/ccj/dataset/plane_ship_tank",
                        help="VOC dataset root directory")
    PARSER.add_argument("--save_dir", type=str,
                        default="/zongshiban_data/2/ccj/dataset/plane_ship_tank_VOC",
                        help="VOC dataset root directory")
    ARGS = PARSER.parse_args()
    ROOT_DIR = ARGS.root_dir
    SAVE_DIR = ARGS.save_dir

    reorganise_directory(ROOT_DIR, SAVE_DIR)